#!/usr/bin/env python
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team
import argparse
import math

import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler

from transformers import (
    AutoModelForCausalLM,
    SchedulerType,
    default_data_collator,
    get_scheduler,
)

import deepspeed
from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
from deepspeed import get_accelerator

from dschat.utils.data.data_utils import create_prompt_dataset
from dschat.utils.utils import print_rank_0, to_device, save_hf_format, set_random_seed, get_all_reduce_mean, get_optimizer_grouped_parameters, save_zero_three_model, load_hf_tokenizer
from dschat.utils.ds_utils import get_train_ds_config
from dschat.utils.module.lora import convert_linear_layer_to_lora, convert_lora_to_linear_layer, only_optimize_lora_parameters, make_model_gradient_checkpointing_compatible
from dschat.utils.model.model_utils import create_hf_model, causal_lm_model_to_fp32_loss
from dschat.utils.perf import print_throughput


def parse_args():
    parser = argparse.ArgumentParser(
        description=
        "Finetune a transformers model on a causal language modeling task")
    # 数据路径。如果是云服务器，则需要先将数据集下载到服务器中。因为国内无法直接从HugFace上自动下载模型和数据。
    parser.add_argument('--data_path',
                        nargs='*',
                        default=['Dahoas/rm-static'],
                        help='Path to the training dataset. Accepted format:'
                        '1) a single data path, 2) multiple datasets in the'
                        'form: dataset1-path dataset2-path ...')
    # 几个阶段的数据划分
    parser.add_argument('--data_split',
                        type=str,
                        default='2,4,4',
                        help='Comma-separated list of proportions for training'
                        'phase 1, 2, and 3 data. For example the split `6,2,2`'
                        'will use 60%% of data for phase 1, 20%% for phase 2'
                        'and 20%% for phase 3.')
    # 只用于SFT（有监督微调，也就是第一步）的数据集路径。如果指定该路径则无需指定data_path，也无需特别划分各个阶段的数据data_split。
    parser.add_argument(
        '--sft_only_data_path',
        nargs='*',
        default=[],
        help='Path to the dataset for only using in SFT phase.')
    # 数据相关文件的存储路径
    parser.add_argument(
        '--data_output_path',
        type=str,
        default='/tmp/data_files/',
        help=
        'Where to store the data-related files such as shuffle index. This needs to be on a local storage of a node (not on a shared storage)'
    )
    # 预训练模型路径-与训练数据同理，下载到云服务器本地。
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        default="/hy-tmp/DeepSpeedExamples/applications/pre/models",
        help=
        "Path to pretrained model or model identifier from huggingface.co/models.",
        required=True,
    )
    # 每块GPU上的训练批次大小  
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=16,
        help="Batch size (per device) for the training dataloader.",
    )
    # 每块GPU验证和测试的批次大小
    parser.add_argument(
        "--per_device_eval_batch_size",
        type=int,
        default=16,
        help="Batch size (per device) for the evaluation dataloader.",
    )
    # 最大序列长度。也就是一个句子最大的token数量。
    parser.add_argument(
        "--max_seq_len",
        type=int,
        default=512,
        help="The maximum sequence length.",
    )
    # 学习率 
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=1e-3,
        help=
        "Initial learning rate (after the potential warmup period) to use.",
    )
    # 权重衰减
    parser.add_argument("--weight_decay",
                        type=float,
                        default=0.,
                        help="Weight decay to use.")
    # 训练轮次
    parser.add_argument("--num_train_epochs",
                        type=int,
                        default=1,
                        help="Total number of training epochs to perform.")
    # 反向传播传播累积：当样本太大时，批次太小时，这种会造成训练不稳定，收敛速度满，训练时间太长。
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help=
        "Number of updates steps to accumulate before performing a backward/update pass.",
    )
    # 调度器类型 动态调整学习率
    parser.add_argument(
        "--lr_scheduler_type",
        type=SchedulerType,
        default="cosine",
        help="The scheduler type to use.",
        choices=[
            "linear", "cosine", "cosine_with_restarts", "polynomial",
            "constant", "constant_with_warmup"
        ],
    )
    # 学习率调度器中的warmup步骤数
    parser.add_argument(
        "--num_warmup_steps",
        type=int,
        default=0,
        help="Number of steps for the warmup in the lr scheduler.")
    # 存储模型路径
    parser.add_argument("--output_dir",
                        type=str,
                        default=None,
                        help="Where to store the model.")
    # 设置随机种子
    parser.add_argument("--seed",
                        type=int,
                        default=1234,
                        help="A seed for reproducible training.")
    # 用于在GPU上进行分布式训练的local rank
    parser.add_argument("--local_rank",
                        type=int,
                        default=-1,
                        help="local_rank for distributed training on gpus")
    # 是否启用梯度的checkpoint
    parser.add_argument('--gradient_checkpointing',
                        action='store_true',
                        help='Enable HF gradient checkpointing for model.')
    # 是否禁用模型的dropout
    parser.add_argument(
        "--dropout",
        type=float,
        default=None,
        help="If dropout configured, use it. "
        "Otherwise, keep the default dropout configuration of the model.")
    # deepspeed features
    # 是否启用Zero offload技术，如果为True，那么模型参数和优化器状态offload到CPU，否则不使用offload
    parser.add_argument('--offload',
                        action='store_true',
                        help='Enable ZeRO Offload techniques.')
    # 混合精度：低精度训练
    parser.add_argument('--dtype',
                        type=str,
                        default='fp16',
                        choices=['fp16', 'bf16'],
                        help='Training data type')
    # ! 用于Actor模型（和clones）的Zero优化阶段
    parser.add_argument(
        '--zero_stage',
        type=int,
        default=0,
        help='ZeRO optimization stage for Actor model (and clones).')
    ## LoRA for efficient training setting
    # 如果大于0，使用LoRA进行高效训练
    parser.add_argument("--lora_dim",
                        type=int,
                        default=0,
                        help="If > 0, use LoRA for efficient training.")
    # Lora的范围
    parser.add_argument("--lora_module_name",
                        type=str,
                        default="decoder.layers.",
                        help="The scope of LoRA.")
    # 只优化LoRA参数
    parser.add_argument('--only_optimize_lora',
                        action='store_true',
                        help='Only optimize the LoRA parameters.')
    # ! Lora的学习率，怎么用? 为什么这么用?
    parser.add_argument(
        "--lora_learning_rate",
        type=float,
        default=5e-4,
        help=
        "Initial LoRA learning rate (after the potential warmup period) to use."
    )
    # ! 损失函数-使用低精度f32计算
    # * 作用: 使用f32计算损失函数。如果使用fp16/bf16会导致精度不足、梯度计算不稳定和损失为
    # * NAN或训练不稳定
    ## low precision
    parser.add_argument(
        '--compute_fp32_loss',
        action='store_true',
        help='Relevant for low precision dtypes (fp16, bf16, etc.). '
        'If specified, loss is calculated in fp32.')
    ## Tensorboard logging
    # 是否启用训练可视化
    parser.add_argument('--enable_tensorboard',
                        action='store_true',
                        help='Enable tensorboard logging')
    # 训练可视化日志保存路径
    parser.add_argument('--tensorboard_path',
                        type=str,
                        default="step1_tensorboard")
    ## Tokenizer
    # 是否添加eot_token，也可以理解为Seq的结束标记
    parser.add_argument(
        "--add_eot_token",
        action='store_true',
        help="Add `eot_token` as additional special token to tokenizer")
    # eot_token：结束文本的特殊Token
    parser.add_argument(
        "--eot_token",
        type=str,
        default="<|endoftext|>",
        help="Specify the format of the `eot_token`",
    )
    ## Print loss
    parser.add_argument('--print_loss',
                        action='store_true',
                        help='Prints loss at each step.')
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()

    return args


def main():
    # 第1步：参数配置
    args = parse_args()

    if args.local_rank == -1:
        device = torch.device(get_accelerator().device_name())
    else:
        get_accelerator().set_device(args.local_rank)
        device = torch.device(get_accelerator().device_name(), args.local_rank)
        # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
        # torch.distributed.init_process_group(backend='nccl')
        # 初始化分布式训练环境
        deepspeed.init_distributed()
    # 获取当前GPU的rank或者索引
    '''
        # * DistributedDataParallel:一个GPU一个进程
        # ! 与 int(os.environ["LOCAL_RANK"])区别 
            1、int(os.environ["LOCAL_RANK"])是获取当前进程在本节点GPU的编号
            2、torch.distributed.get_rank()是当前进程在所有进程中的全局编号
    '''
    args.global_rank = torch.distributed.get_rank()
    # 根据输入参数返回一个训练数据集的配置字典
    ds_config = get_train_ds_config(offload=args.offload,
                                    dtype=args.dtype,
                                    stage=args.zero_stage,
                                    enable_tensorboard=args.enable_tensorboard,
                                    tb_path=args.tensorboard_path,
                                    tb_name="step1_model")
    ds_config[
        'train_micro_batch_size_per_gpu'] = args.per_device_train_batch_size
    ds_config[
        'train_batch_size'] = args.per_device_train_batch_size * torch.distributed.get_world_size(
        ) * args.gradient_accumulation_steps

    # If passed along, set the training seed now.
    set_random_seed(args.seed)
    # ! 分布式训练的同步工具，确保所有的分布式进程都达到这个阻塞点
    # * 防止快速加载好的GPU提前进入数据集分配环节，影响整个全局进程。
    torch.distributed.barrier()

    # load_hf_tokenizer will get the correct tokenizer and set padding tokens based on the model family
    # 加载token字典。
    additional_special_tokens = args.eot_token if args.add_eot_token else None
    tokenizer = load_hf_tokenizer(args.model_name_or_path,
                                  fast_tokenizer=True,
                                  add_special_tokens=additional_special_tokens)
    # 第2步：创建actor模型
    model = create_hf_model(AutoModelForCausalLM,
                            args.model_name_or_path,
                            tokenizer,
                            ds_config,
                            dropout=args.dropout)

    if args.compute_fp32_loss:
        print_rank_0(
            f"Using model {model.__class__.__name__} with loss in fp32",
            args.global_rank)
        causal_lm_model_to_fp32_loss(model)

    if args.lora_dim > 0:
        # ! 将模型指定的线性层，增加Lora分枝层
        model = convert_linear_layer_to_lora(model, args.lora_module_name,
                                             args.lora_dim)
        # 判定是否只优化lora层
        if args.only_optimize_lora:
            # 只优化lora层，则需要冻结其他层参数
            model = only_optimize_lora_parameters(model)
            # ! ?作用是什么 如何理解
            # * 回答：如果只优化Lora层，则embedding层输出值的属性没有要求计算梯度（requires_grad=False）。那么，想要优化lora层，则需要将embedding层的输出requires_grad设置为True
            model = make_model_gradient_checkpointing_compatible(model)

    # Prepare the data
    # 第3步 准备数据集
    train_dataset, eval_dataset = create_prompt_dataset(
        args.local_rank,
        args.data_path,
        args.data_split,
        args.data_output_path,
        train_phase,
        args.seed,
        tokenizer,
        args.max_seq_len,
        end_of_conversation_token=tokenizer.eos_token,
        sft_only_data_path=args.sft_only_data_path) # 专用于有监督微调的数据集
    # DataLoaders creation:
    # 不在分布式训练环境下，因此我们将使用随机采样和顺序采样
    if args.local_rank == -1:
        # 在训练过程，随机选择样本进行训练，防止模型过拟合
        train_sampler = RandomSampler(train_dataset)
        # 在评估过程，将按照数据集中的顺序选择样本进行评估，按照顺序不会影响模型的性能评估
        eval_sampler = SequentialSampler(eval_dataset)
    # 在分布式训练环境中，将使用分布式采样
    else:
        # 创建一个用于训练集的分布式采样器，会在所有的训练节点上对样本进行均匀的分布采样
        # 确保每个节点处理的样本是独立且均匀的，从而提高分布式训练的效率和稳定性
        train_sampler = DistributedSampler(train_dataset)
        # 创建一个用于评估集的分布式采样器
        eval_sampler = DistributedSampler(eval_dataset)
    # 创建用于训练集的数据加载器
    train_dataloader = DataLoader(train_dataset,
                                  collate_fn=default_data_collator,
                                  sampler=train_sampler,
                                  batch_size=args.per_device_train_batch_size)
    # 创建用于评估集的数据加载器
    eval_dataloader = DataLoader(eval_dataset,
                                 collate_fn=default_data_collator,
                                 sampler=eval_sampler,
                                 batch_size=args.per_device_eval_batch_size)
    # 模型评估
    def evaluation(model, eval_dataloader):
        # 将模型设置为
        model.eval()
        losses = 0
        for step, batch in enumerate(eval_dataloader):
            batch = to_device(batch, device)
            with torch.no_grad():
                # 用当前批次的输入数据去前向传播模型，并得到模型的输出
                outputs = model(**batch)
            # 提取模型输出的损失值
            loss = outputs.loss
            # 累积整个评估过程中的损失值
            losses += loss.float()
        # 计算平均损失值
        losses = losses / (step + 1)
        # ! 和原来版本的差异性：原版本是先计算困惑度然后计算困惑度的均值
        try:
            # 若果是一个分布式设置，该函数会计算所有device上的平均损失
            losses = get_all_reduce_mean(losses)
        except:
            pass
        try:
            # 计算模型的困惑度，这是评估语言模型性能的常用指标。困惑度的计算方法是对平均损失值的取指数
            # 如果这步运算没有发生溢出，那么困惑度的值就是torch.exp(losses)
            # 当损失值过大时，指数运算可能会导致溢出
            perplexity = torch.exp(losses).item()
        except OverflowError:
            # 如果上一步的指数运算发生溢出，那么将困惑度的值设为无穷大
            perplexity = float("inf")
        return perplexity, losses.item()
    
    # 第4步 将模型参数分组、创建优化器 和 学习率调度器
    # Split weights in two groups, one with weight decay and the other not.
    # 1、将模型参数分为三组，一组源模型应用权重衰减，一组是lora层的权重衰减，另一组不应用
    optimizer_grouped_parameters = get_optimizer_grouped_parameters(
        model, args.weight_decay, args.lora_learning_rate)
    # 根据是否使用DeepSpeed的CPU offload功能来选择优化器，优化器定义了如何更新模型的参数以最小化损失函数
    # DeepSpeedCPUAdam : 为了配合DeepSpeed的CPU offload功能设计的，在DeepSpeed中，CPU offload可以将模型参数、优化器状态和梯度数据在CPU和GPU之间进行切换，以减轻GPU的内存压力。
    # FuseAdam ： 将一些计算操作融合在一起(fused)，以减少计算时间和内存使用量。FusedAdam主要是为了提高在GPU上的运算效率。
    AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
    # 使用选择的优化器进行实例化
    optimizer = AdamOptimizer(optimizer_grouped_parameters,
                              lr=args.learning_rate,
                              betas=(0.9, 0.95))
    # 计算每个epoch需要进行的更新步数，等于训练数据集大小除以梯度累积步数（对结果向上取整）
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps)
    # 创建学习率调度器
    lr_scheduler = get_scheduler(
        name=args.lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=args.num_warmup_steps, # ? 学习率预热，让学习率先从很小的值逐步提升到目标学习率。在这之后调用args.lr_scheduler_type往下调整学习率
        num_training_steps=args.num_train_epochs * num_update_steps_per_epoch, # 总的训练步数等于训练的epoch数乘以每个epoch的更新步数
    )
    # 第5步：deepspeed初始化，创建模型、优化器、学习率调度器
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
        model=model,
        optimizer=optimizer,
        args=args,
        config=ds_config, # deepspeed的配置信息
        lr_scheduler=lr_scheduler,
        dist_init_required=True # 需要进行分布式训练的初始化
    )
    # 是否启用梯度检查点
    if args.gradient_checkpointing:
        # 梯度检查点是一种可以减少训练过程中内存使用的技术，通过牺牲一部分计算效果来换取内存占用的降低
        # * 作用: HuggingFace的gradient_checkpointing_enable()函数，自动将部分层或者模块设置为检查点，而不是所有layer设置为检查点。eg:TransFormer的 block
        model.gradient_checkpointing_enable()
    # 第6步：模型的训练和验证
    # Train!
    print_rank_0("***** Running training *****", args.global_rank)
    print_rank_0(
        f"***** Evaluating perplexity, Epoch {0}/{args.num_train_epochs} *****",
        args.global_rank)
    # 模型评估    
    perplexity, eval_loss = evaluation(model, eval_dataloader)
    print_rank_0(f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank)
    # 模型训练
    for epoch in range(args.num_train_epochs):
        print_rank_0(
            f"Beginning of Epoch {epoch+1}/{args.num_train_epochs}, Total Micro Batches {len(train_dataloader)}",
            args.global_rank)
        model.train()
        import time
        for step, batch in enumerate(train_dataloader):
            start = time.time()
            # 将数据移动到指定的训练设备
            batch = to_device(batch, device)
            outputs = model(**batch, use_cache=False)
            loss = outputs.loss
            if args.print_loss:
                print(
                    f"Epoch: {epoch}, Step: {step}, Rank: {torch.distributed.get_rank()}, loss = {loss}"
                )
            model.backward(loss)
            model.step()
            end = time.time()
            if torch.distributed.get_rank() == 0:
                print_throughput(model.model, args, end - start,
                                 args.global_rank)

        # Evaluate perplexity on the validation set.
        print_rank_0(
            f"***** Evaluating perplexity, Epoch {epoch+1}/{args.num_train_epochs} *****",
            args.global_rank)
        # 验证集推理
        perplexity, eval_loss = evaluation(model, eval_dataloader)
        print_rank_0(f"ppl: {perplexity}, loss: {eval_loss}", args.global_rank)
        # 更新了模型的内部计时器，表示一个epoch已经完成
        model.tput_timer.update_epoch_count()
    # 第7步 保存模型
    if args.output_dir is not None:
        print_rank_0('saving the final model ...', args.global_rank)
        # 将模型中的Lora层转换为标准的线性层， 这样使得模型在保存后可以在你没有Lora层代码的环境中加载和使用。实际操作：权重融合，并为转换
        model = convert_lora_to_linear_layer(model)
        # 是否在主进程
        if args.global_rank == 0:
            # 保存模型
            save_hf_format(model, tokenizer, args)
        # 是否使用Zero Redundancy Optimizer的第三个阶段（Zero-3）
        # Zero-3是一种内存优化策略，可以大大减少模型训练中所需的GPU内存，但同时也意味着模型的各部分放在不同的GPU中
        if args.zero_stage == 3:
            # For zero stage 3, each gpu only has a part of the model, so we need a special save function
            save_zero_three_model(model,
                                  args.global_rank,
                                  args.output_dir,
                                  zero_stage=args.zero_stage)


if __name__ == "__main__":
    main()
